import argparse  # 导入参数库
import datetime
import queue  # 导入队列库
import socket
import threading  # 导入线程库

import prettytable  # 导入表格库
import pyfiglet  # 导入字体库
from colorama import Fore  # 导入颜色库


# 横幅
def banner():
    print(Fore.BLUE + "PortScan v1.0  简单的端口扫描工具   Author: Youze")
    test = pyfiglet.figlet_format("PortScan", font="calvin_s")
    print(Fore.BLUE + test)


# 命令参数获取
def get_args():
    parser = argparse.ArgumentParser(description="端口扫描工具")
    parser.add_argument(
        "-u", "--target", help="目标地址", required=True
    )  # required=True表示必须输入
    parser.add_argument(
        "-v", "--version", action="version", version="%(prog)s 2.0", help="显示版本信息"
    )
    parser.add_argument(
        "-p",
        "--port",
        help="扫描端口范围 all 或 normal 或 start-end 或 port1,port2,port3",
        default="normal",
    )
    parser.add_argument("-o", "--output", help="输出文件", default="")
    parser.add_argument("-n", "--number", help="线程数量", default=10, type=int)
    return parser.parse_args()


# 端口格式检查
def port_check(port):
    if port == "normal":
        default_ports = [
            20,
            21,
            22,
            23,
            25,
            53,
            80,
            110,
            143,
            443,
            445,
            993,
            995,
            1723,
            3306,
            3389,
            5432,
            8080,
            8443,
            1521,
            5900,
            6379,
            10000,
            33000,
        ]
        return default_ports
    elif port == "all":
        return list(range(1, 65536))  # 65535 + 1
    elif "-" in port:
        start_port, end_port = map(int, port.split("-"))
        if start_port < 1 or start_port > 65535 or end_port < 1 or end_port > 65535:
            print(Fore.RED + "端口范围错误")
            exit()
        elif start_port > end_port:
            start_port, end_port = end_port, start_port  # 交换
        return list(range(start_port, end_port + 1))
    else:
        return [int(p) for p in port.split(",")]


queue = queue.Queue()  # 队列
threads = []  # 线程列表
results = []  # 扫描结果列表
table = prettytable.PrettyTable(["Port", "Time"])  # 表格


def scan_port(host, queue):
    while not queue.empty():  # 队列不为空时循环
        port = queue.get()
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(1)  # 设置超时
        try:
            result = sock.connect_ex((host, port))
            time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")  # 格式化时间
            if result == 0:
                results.append((port, time))
        except Exception as e:
            print(Fore.RED + f"错误: {e}")
        finally:
            sock.close()  # 确保 socket 被关闭


def thread_scan(srgs_host, srgs_port, srgs_output, srgs_number):
    # 将srgs_port中的端口放入队列
    for p in srgs_port:
        queue.put(p)
    # 创建线程
    for i in range(srgs_number):
        t = threading.Thread(target=scan_port, args=(srgs_host, queue))
        threads.append(t)
        t.start()

    # 等待所有线程完成
    for t in threads:
        t.join()

    # 添加结果到表格
    for port, time in results:
        table.add_row([port, time])

    # 输出结果
    if srgs_output:
        with open(srgs_output, "w") as f:
            f.write(str(table))
        print(Fore.GREEN + f"结果已保存到 {srgs_output}")

    print(table)


if __name__ == "__main__":
    banner()
    srgs = get_args()
    srgs_host = srgs.target
    srgs_port = port_check(srgs.port)
    srgs_output = srgs.output
    srgs_number = srgs.number
    print(Fore.GREEN + "[+]开始扫描")
    thread_scan(srgs_host, srgs_port, srgs_output, srgs_number)
